#include <dirent.h>
#include <gsl/gsl_matrix_double.h>
#include <string.h>

#include "solver.h"
#include "global_definitions.h"
#include "math.h"
#include "stdio.h"
#include "stdlib.h"
#include "omp.h"
#include "lib.h"
#include "initialize_config.h"
/*
 * DECLATION OF HELPER FUNCTONS
 * 
 */
double get_pulse_in_zt ( fsp_t*, double, double );
double get_pulse_dz_in_zt ( fsp_t*, double, double );
double get_pulse_dt_in_zt ( fsp_t*, double, double );

int f ( double, const double*, double*, void * );
int jacobian ( double, const double*, double *, double *, void * );
int solver_simple_iteration ( system_paramaters_t*,
                              double start_t, double stop_t, double* );
int one_elementary_calc_loop ( system_paramaters_t* system_parameters, element_of_solution_t* element_of_solution );

int check_electron ( double v_z, double t0, double z0, double t1, double z1, double dz );
int reset_element_of_solution ( element_of_solution_t* );
int initialize_element_of_solution ( element_of_solution_t*, fsp_t*, double z0, int nz0, int i_tp, int i_zp );
int save_result_into_elementary_solution ( element_of_solution_t*, double*, double, double, double, int );

/*
 *DEFINITIONS OF INTERFACE FUNCTIONS
 */

int solver_initialize_ode_system ( system_paramaters_t* system_paramaters, fsp_t* fsp )
{
    gsl_odeiv_system* ode_system = (gsl_odeiv_system*) malloc( sizeof (gsl_odeiv_system) );
    const gsl_odeiv_step_type * step_type = gsl_odeiv_step_rkf45;

    printf( "Initializing ODE System\n" );
    system_paramaters->fsp = fsp;
    system_paramaters->ode_system = ode_system;

    ode_system->function = &f;
    ode_system->jacobian = &jacobian;
    ode_system->dimension = SYSTEM_DIMENSION;
    ode_system->params = (void*) fsp;

    system_paramaters->ode_step = gsl_odeiv_step_alloc( step_type, SYSTEM_DIMENSION );
    system_paramaters->ode_control = gsl_odeiv_control_y_new( fsp->simp->abs_eps, fsp->simp->rel_eps );
    system_paramaters->ode_evolve = gsl_odeiv_evolve_alloc( SYSTEM_DIMENSION );
    printf( "ODE System is initialized\n" );
    return GSL_SUCCESS;
}

int solver_free_resources_for_ode_system ( system_paramaters_t* system_paramaters )
{
    printf( "Free ODE System resources\n" );
    gsl_odeiv_step_free( system_paramaters->ode_step );
    gsl_odeiv_control_free( system_paramaters->ode_control );
    gsl_odeiv_evolve_free( system_paramaters->ode_evolve );
    free( system_paramaters->ode_system );
    printf( "ODE System is unitialized\n" );
    return GSL_SUCCESS;
}

int solver_get_full_solution ( system_paramaters_t* system_parameters )
{
    fsp_t* fsp = system_parameters->fsp;
    simp_t* simp = fsp->simp;
    csimp_t* csimp = fsp->csimp;
    double max_abs_v_z = MAX( gsl_matrix_max( csimp->zsd ), fabs( gsl_matrix_min( csimp->zsd ) ) );
    double max_t_point = gsl_vector_max( simp->time_collect_points );
    double l = max_abs_v_z*max_t_point;
    int i_tp = 0;
    int i_zp = 0;
    double z0 = 0;
    double dz = simp->dz;
    int nz0 = 0;
    int num_of_t_points = simp->time_collect_points->size;
    int num_of_z_points = simp->z_collect_points->size;
    char* folder_to_save = simp->savedir;
    char* name_template = simp->name_template;

    for ( z0 = -l; z0 <= l; z0 += dz )
    {

        for ( i_tp = 0; i_tp < num_of_t_points; ++i_tp )
        {
#pragma omp parallel for
            for ( i_zp = 0; i_zp < num_of_z_points; ++i_zp )
            {
                int num_of_calc_particles_in_elementary_sol = 0;
                element_of_solution_t element_of_solution;

                initialize_element_of_solution( &element_of_solution, fsp, z0, nz0, i_tp, i_zp );
                num_of_calc_particles_in_elementary_sol =
                        one_elementary_calc_loop( system_parameters,
                                                  &element_of_solution );

                if ( num_of_calc_particles_in_elementary_sol > 0 )
                {
                    element_of_solution.num_of_rows = num_of_calc_particles_in_elementary_sol;
                    solver_save_element_of_solution_into_file( folder_to_save,
                                                               name_template,
                                                              &element_of_solution );
                }
                solver_free_element_of_solution( &element_of_solution );
            }
        }
        ++nz0;
    }
}

int solver_read_element_of_solution_from_file(char* filename, element_of_solution_t* element_of_solution)
{
    FILE* file = fopen(filename, "r");
    fread( &(element_of_solution->t), 1, sizeof (element_of_solution->t), file );
    fread( &(element_of_solution->z), 1, sizeof (element_of_solution->z), file );
    fread( &(element_of_solution->z0), 1, sizeof (element_of_solution->z0), file );
    fread( &(element_of_solution->nt), 1, sizeof (element_of_solution->nt), file );
    fread( &(element_of_solution->nz), 1, sizeof (element_of_solution->nz), file );
    fread( &(element_of_solution->nz0), 1, sizeof (element_of_solution->nz0), file );
    fread( &(element_of_solution->num_of_rows), 1, sizeof (element_of_solution->num_of_rows), file );

    element_of_solution->stat = gsl_matrix_alloc(element_of_solution->num_of_rows, 2);
    gsl_matrix_fread(file, element_of_solution->stat);
    return GSL_SUCCESS;
    
}

int solver_save_element_of_solution_into_file (char* dir, char* name_template,
                                               element_of_solution_t* element_of_solution)
{
    char filename[MAXNAMLEN];
    FILE* file;
    int nz0 = element_of_solution->nz0;
    int ntp = element_of_solution->nt;
    int nzp = element_of_solution->nz;
    sprintf( filename, "%s/%s_%d_%d_%d.dat", dir, name_template, nz0, ntp, nzp );
    file = fopen( filename, "w" );

    fwrite( &(element_of_solution->t), 1, sizeof (element_of_solution->t), file );
    fwrite( &(element_of_solution->z), 1, sizeof (element_of_solution->z), file );
    fwrite( &(element_of_solution->z0), 1, sizeof (element_of_solution->z0), file );
    fwrite( &(element_of_solution->nt), 1, sizeof (element_of_solution->nt), file );
    fwrite( &(element_of_solution->nz), 1, sizeof (element_of_solution->nz), file );
    fwrite( &(element_of_solution->nz0), 1, sizeof (element_of_solution->nz0), file );
    fwrite( &(element_of_solution->num_of_rows), 1, sizeof (element_of_solution->num_of_rows), file );
    gsl_matrix_view only_nonzero_particles = gsl_matrix_submatrix( element_of_solution->stat,
                                                                   0, 0, element_of_solution->num_of_rows,
                                                                   2 );
    gsl_matrix_fwrite( file, &only_nonzero_particles.matrix );
    fclose( file );
    return GSL_SUCCESS;
}

/*
 * DEFINITIONS OF HELPER FUNCTIONS
 */

int f ( double t, const double y[], double dydt[], void * params )
{
    /*
     * y[0]..y[2] --- v_x, v_y, v_z, correspondigly
     * y[3]..y[5] --- x, y, z, correspondigly
     */
    static double abs_e_m = -(ELECTRON_CHARGE / ELECTRON_MASS);
    fsp_t* fsp = (fsp_t*) params;
    double E_x = get_pulse_in_zt( fsp, t, y[5] );
    double w_ce = fsp->csimp->w_ce;

    dydt[0] = -abs_e_m * E_x - w_ce * y[1]; //d(v_x)dt
    dydt[1] = w_ce * y[0]; //d(v_y)dt
    dydt[2] = 0; //d(v_z)dt
    dydt[3] = 0; //dxdt
    dydt[4] = 0; //dydt
    dydt[5] = y[2]; //dzdt

    return GSL_SUCCESS;
}

int jacobian ( double t, const double y[], double * dfdy, double dfdt[], void * params )
{
    fsp_t* fsp = (fsp_t*) params;
    double w_ce = fsp->csimp->w_ce;
    double JE_xdz = get_pulse_dz_in_zt( fsp, t, y[5] );
    double JE_xdt = get_pulse_dt_in_zt( fsp, t, y[5] );
    gsl_matrix_view dfdy_mat = gsl_matrix_view_array( dfdy, 6, 6 );
    gsl_matrix * m = &dfdy_mat.matrix;

    gsl_matrix_set_all( m, 0 );

    gsl_matrix_set( m, 0, 1, -w_ce );
    gsl_matrix_set( m, 0, 5, JE_xdz );

    gsl_matrix_set( m, 1, 0, w_ce );

    gsl_matrix_set( m, 5, 2, 1 );

    dfdt[0] = JE_xdt;
    dfdt[1] = 0;
    dfdt[2] = 0;
    dfdt[3] = 0;
    dfdt[4] = 0;
    dfdt[5] = 0;

    return GSL_SUCCESS;
}

double get_pulse_in_zt ( fsp_t* fsp, double t, double z )
{
    double dz = fsp->simp->dz;
    double dt = fsp->csimp->dt;
    int row = lib_gsl_bsearch( fsp->csimp->pulse_time, t, dt, 0, fsp->csimp->pulse_time->size - 1 );
    int col = lib_gsl_bsearch( fsp->csimp->pulse_z, z, dz, 0, fsp->csimp->pulse_z->size - 1 );
    if ( (row > 0) && (col > 0) )
    {
        return gsl_matrix_get( fsp->csimp->pulse, row, col );
    } else
    {
        return 0;
    }

}

double get_pulse_dt_in_zt ( fsp_t* fsp, double t, double z )
{
    double dz = fsp->simp->dz;
    double dt = fsp->csimp->dt;
    int row = lib_gsl_bsearch( fsp->csimp->pulse_time, t, dt, 0, fsp->csimp->pulse_time->size - 1 );
    int col = lib_gsl_bsearch( fsp->csimp->pulse_z, z, dz, 0, fsp->csimp->pulse_z->size - 1 );
    if ( (row > 0) && (col > 0) )
    {
        return gsl_matrix_get( fsp->csimp->pulse_dt, row, col );
    } else
    {
        return 0;
    }

}

double get_pulse_dz_in_zt ( fsp_t* fsp, double t, double z )
{
    double dz = fsp->simp->dz;
    double dt = fsp->csimp->dt;
    int row = lib_gsl_bsearch( fsp->csimp->pulse_time, t, dt, 0, fsp->csimp->pulse_time->size - 1 );
    int col = lib_gsl_bsearch( fsp->csimp->pulse_z, z, dz, 0, fsp->csimp->pulse_z->size - 1 );
    if ( (row > 0) && (col > 0) )
    {
        return gsl_matrix_get( fsp->csimp->pulse_dz, row, col );
    } else
    {
        return 0;
    }

}

int check_electron ( double v_z, double t0, double z0, double t1, double z1, double half_dz )
{
    double e_path = v_z * (t1 - t0);
    double Dz = z1 - z0;

    if ( fabs( Dz - e_path ) <= half_dz )
    {
        return 1;
    } else
    {
        return 0;
    }

}

int solver_simple_iteration ( system_paramaters_t* system_paramaters,
                              double start_t, double stop_t, double y[] )
{
    //FILE* fsave = fopen(filename, "w");
    double h = system_paramaters->fsp->csimp->dt;
    double t = start_t;
    double t1 = stop_t;
    int status = GSL_FAILURE;

    gsl_odeiv_evolve_reset( system_paramaters->ode_evolve );
    while ( t < t1 )
    {
        status = gsl_odeiv_evolve_apply( system_paramaters->ode_evolve,
                                         system_paramaters->ode_control,
                                         system_paramaters->ode_step,
                                         system_paramaters->ode_system,
                                         &t, t1,
                                         &h, y );
        if ( status != GSL_SUCCESS )
        {
            break;
        }
    }
    return status;
}

int reset_element_of_solution ( element_of_solution_t* element_of_solution )
{
    element_of_solution->t = 0;
    element_of_solution->z = 0;
    gsl_matrix_set_zero( element_of_solution->stat );
    return GSL_SUCCESS;
}

int initialize_element_of_solution ( element_of_solution_t* element_of_solution,
                                     fsp_t* fsp, double z0, int nz0, int i_tp, int i_zp )
{
    int fullsize = (fsp->simp->NumOfZElectrons)*(fsp->simp->NumOfTrElectrons)*(fsp->simp->NumOfTrElectrons);
    //int fullsize = (fsp->simp->NumOfTrElectrons)*(fsp->simp->NumOfTrElectrons);
    element_of_solution->stat = gsl_matrix_alloc( fullsize, 2 );
    element_of_solution->nt = i_tp;
    element_of_solution->nz = i_zp;

    element_of_solution->t = gsl_vector_get( fsp->simp->time_collect_points, i_tp );

    element_of_solution->z = gsl_vector_get( fsp->simp->z_collect_points, i_zp );

    element_of_solution->z0 = z0;
    element_of_solution->nz0 = nz0;

    return GSL_SUCCESS;
}

int solver_free_element_of_solution ( element_of_solution_t* element_of_solution )
{
    gsl_matrix_free( element_of_solution->stat );
}

int save_result_into_elementary_solution ( element_of_solution_t* element_of_solution,
                                             double* y, double xs_pdf, double ys_pdf, double zs_pdf,
                                             int row_to_save )
{
    double v_x = y[ELEMENT_OF_SOLUTION_XSPEED_VAL];
    double v_y = y[ELEMENT_OF_SOLUTION_YSPEED_VAL];
    double v_tr = sqrt(v_x*v_x + v_y*v_y);
    gsl_matrix_set( element_of_solution->stat, row_to_save,
                    ELEMENT_OF_SOLUTION_V_TR_VAL, v_tr );

    /*gsl_matrix_set( element_of_solution->stat, row_to_save,
                    ELEMENT_OF_SOLUTION_YSPEED_VAL, y[ELEMENT_OF_SOLUTION_YSPEED_VAL] );

    gsl_matrix_set( element_of_solution->stat, row_to_save,
                    ELEMENT_OF_SOLUTION_ZSPEED_VAL, y[ELEMENT_OF_SOLUTION_ZSPEED_VAL] );*/

    /*gsl_matrix_set( element_of_solution->stat, row_to_save, ELEMENT_OF_SOLUTION_XSPEED_PDF, xs_pdf );
    gsl_matrix_set( element_of_solution->stat, row_to_save, ELEMENT_OF_SOLUTION_YSPEED_PDF, ys_pdf );*/
    gsl_matrix_set( element_of_solution->stat, row_to_save, ELEMENT_OF_SOLUTION_V_TR_PDF, xs_pdf*ys_pdf);
    return GSL_SUCCESS;
}

int one_elementary_calc_loop ( system_paramaters_t* system_parameters,  element_of_solution_t* element_of_solution )
{
    int i_zel = 0;
    simp_t* simp = system_parameters->fsp->simp;
    csimp_t* csimp = system_parameters->fsp->csimp;
    double v_z0 = 0;
    double z0 = element_of_solution->z0;
    double y0[6] = {0.0, 0.0, 0.0, 0.0, 0.0, z0};
    double y[6] = {0.0, 0.0, 0.0, 0.0, 0.0, 0.0};
    double t0 = 0;
    double t1 = 0;
    double z1 = 0;
    double zs_pdf = 0;
    double xs_pdf = 0;
    double ys_pdf = 0;
    int i_trelx = 0;
    int i_trely = 0;
    double half_dz = 0.5 * simp->dz;
    int num_of_calc_particles_in_elementary_sol = 0;
    int i = 0;


    for ( i_zel = 0; i_zel < simp->NumOfZElectrons; ++i_zel )
    {
        v_z0 = gsl_matrix_get( csimp->zsd, i_zel, COLUMN_ZSPEED_VALUE );
        y0[ELEMENT_OF_SOLUTION_ZSPEED_VAL] = v_z0;
        zs_pdf = gsl_matrix_get( csimp->zsd, i_zel, COLUMN_ZSPEED_PDF );
        t0 = 0;
        t1 = element_of_solution->t;
        z1 = element_of_solution->z;
        if ( check_electron( v_z0, t0, z0, t1, z1, half_dz ) )
        {
            for ( i_trelx = 0; i_trelx < simp->NumOfTrElectrons; ++i_trelx )
            {
                y0[ELEMENT_OF_SOLUTION_XSPEED_VAL] = gsl_matrix_get( csimp->trsd, i_trelx, COLUMN_XSPEED_VALUE );
                xs_pdf = gsl_matrix_get( csimp->trsd, i_trelx, COLUMN_XSPEED_PDF );

                for ( i_trely = 0; i_trely < simp->NumOfTrElectrons; ++i_trely )
                {

                    y0[ELEMENT_OF_SOLUTION_YSPEED_VAL] = gsl_matrix_get( csimp->trsd, i_trely, COLUMN_YSPEED_VALUE );

                   
                    solver_simple_iteration( system_parameters, t0, t1, y0 );


                    ys_pdf = gsl_matrix_get( csimp->trsd, i_trely, COLUMN_YSPEED_PDF );

                    save_result_into_elementary_solution( element_of_solution,
                                                            y0, xs_pdf, ys_pdf, zs_pdf,
                                                            num_of_calc_particles_in_elementary_sol );

                    num_of_calc_particles_in_elementary_sol++;
                }

            }


        }

    }
    return num_of_calc_particles_in_elementary_sol;
}